# -*- coding: utf-8 -*-
from openchemie import OpenChemIE
from PIL import Image
import os

def extract_molecules_from_image():
    """
    演示如何从单个图片文件中识别和提取分子信息。
    """
    print("Initializing OpenChemIE...")
    # 初始化模型，第一次运行时会自动下载所需模型
    model = OpenChemIE()
    print("Initialization complete.")

    # 支持多种图片格式，例如 .png, .jpg, .tif
    image_path = 'example/img1.png'
    
    if not os.path.exists(image_path):
        print(f"Error: Example image not found at '{image_path}'")
        return
        
    print(f"\nProcessing image: {image_path}")
    
    # 使用 PIL (Pillow) 库打开图片
    # OpenChemIE 同时支持 PIL 和 OpenCV 图像对象
    img = Image.open(image_path)
    images = [img]

    # 调用核心方法从图片列表中提取分子信息
    results = model.extract_molecules_from_figures(images)

    # 提取结果是一个列表，对应输入的图片列表
    # 我们检查第一张图片的结果中是否包含分子信息
    if results and results[0].get('molecules'):
        # 获取第一个被识别到的分子
        first_molecule = results[0]['molecules'][0]
        smiles_string = first_molecule.get('smiles', 'N/A')
        
        print("\n--- Extraction Result ---")
        print(f"Detected SMILES from {image_path}: {smiles_string}")
        print("-------------------------")
    else:
        print(f"No molecules were detected in {image_path}.")

if __name__ == "__main__":
    extract_molecules_from_image()
